"""
数据预处理模块
负责数据清洗、特征工程和数据标准化
"""

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
import warnings
warnings.filterwarnings('ignore')

class DataPreprocessor:
    """数据预处理器"""
    
    def __init__(self):
        self.scaler = StandardScaler()
        self.label_encoders = {}
        self.feature_columns = []
        self.is_fitted = False
        
    def preprocess_user_input(self, user_data):
        """
        预处理用户输入数据
        
        Args:
            user_data (dict): 用户输入的数据
            
        Returns:
            dict: 预处理后的数据
        """
        try:
            # 基础信息处理
            processed_data = {}
            
            # 年龄处理
            age = user_data.get('age', 50)
            processed_data['age'] = float(age)
            
            # BMI处理
            height = user_data.get('height', 170)  # cm
            weight = user_data.get('weight', 70)   # kg
            bmi = weight / ((height / 100) ** 2)
            processed_data['bmi'] = bmi
            
            # 性别处理 (0: 女, 1: 男)
            gender = user_data.get('gender', 'male')
            processed_data['gender'] = 1 if gender == 'male' else 0
            
            # 家族史处理
            family_history = user_data.get('family_history', False)
            processed_data['family_history'] = 1 if family_history else 0
            
            # 吸烟状态
            smoking = user_data.get('smoking', False)
            processed_data['smoking'] = 1 if smoking else 0
            
            # 饮酒状态
            drinking = user_data.get('drinking', False)
            processed_data['drinking'] = 1 if drinking else 0
            
            # 运动频率 (0: 从不, 1: 偶尔, 2: 经常, 3: 每天)
            exercise_freq = user_data.get('exercise_frequency', 'occasionally')
            exercise_map = {'never': 0, 'occasionally': 1, 'regularly': 2, 'daily': 3}
            processed_data['exercise_frequency'] = exercise_map.get(exercise_freq, 1)
            
            # 饮食偏好 (0: 清淡, 1: 正常, 2: 重口味)
            diet_preference = user_data.get('diet_preference', 'normal')
            diet_map = {'light': 0, 'normal': 1, 'heavy': 2}
            processed_data['diet_preference'] = diet_map.get(diet_preference, 1)
            
            # 血压数据
            systolic_bp = user_data.get('systolic_bp', 120)
            diastolic_bp = user_data.get('diastolic_bp', 80)
            processed_data['systolic_bp'] = float(systolic_bp)
            processed_data['diastolic_bp'] = float(diastolic_bp)
            
            # 血糖数据
            blood_glucose = user_data.get('blood_glucose', 100)
            processed_data['blood_glucose'] = float(blood_glucose)
            
            # 胆固醇数据
            cholesterol = user_data.get('cholesterol', 200)
            processed_data['cholesterol'] = float(cholesterol)
            
            # 心率数据
            heart_rate = user_data.get('heart_rate', 72)
            processed_data['heart_rate'] = float(heart_rate)
            
            # 计算衍生特征
            processed_data = self._calculate_derived_features(processed_data)
            
            return processed_data
            
        except Exception as e:
            print(f"用户数据预处理错误: {e}")
            return {}
    
    def preprocess_batch_data(self, df):
        """
        批量数据预处理
        
        Args:
            df (DataFrame): 原始数据框
            
        Returns:
            DataFrame: 预处理后的数据框
        """
        try:
            # 复制数据框
            processed_df = df.copy()
            
            # 处理缺失值
            processed_df = self._handle_missing_values(processed_df)
            
            # 处理异常值
            processed_df = self._handle_outliers(processed_df)
            
            # 特征工程
            processed_df = self._feature_engineering(processed_df)
            
            # 数据标准化
            if not self.is_fitted:
                self._fit_scaler(processed_df)
            
            # 应用标准化
            numeric_columns = processed_df.select_dtypes(include=[np.number]).columns
            processed_df[numeric_columns] = self.scaler.transform(processed_df[numeric_columns])
            
            return processed_df
            
        except Exception as e:
            print(f"批量数据预处理错误: {e}")
            return df
    
    def _handle_missing_values(self, df):
        """处理缺失值"""
        # 数值型列用中位数填充
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        for col in numeric_columns:
            if df[col].isnull().any():
                df[col].fillna(df[col].median(), inplace=True)
        
        # 分类型列用众数填充
        categorical_columns = df.select_dtypes(include=['object']).columns
        for col in categorical_columns:
            if df[col].isnull().any():
                df[col].fillna(df[col].mode()[0], inplace=True)
        
        return df
    
    def _handle_outliers(self, df):
        """处理异常值"""
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        
        for col in numeric_columns:
            if col in ['age', 'bmi', 'systolic_bp', 'diastolic_bp', 'blood_glucose']:
                # 使用IQR方法处理异常值
                Q1 = df[col].quantile(0.25)
                Q3 = df[col].quantile(0.75)
                IQR = Q3 - Q1
                lower_bound = Q1 - 1.5 * IQR
                upper_bound = Q3 + 1.5 * IQR
                
                # 将异常值限制在边界内
                df[col] = df[col].clip(lower=lower_bound, upper=upper_bound)
        
        return df
    
    def _feature_engineering(self, df):
        """特征工程"""
        # BMI分类
        df['bmi_category'] = pd.cut(df['bmi'], 
                                   bins=[0, 18.5, 25, 30, 100], 
                                   labels=[0, 1, 2, 3])  # 0: 偏瘦, 1: 正常, 2: 超重, 3: 肥胖
        
        # 血压分类
        df['bp_category'] = 0  # 正常
        df.loc[(df['systolic_bp'] >= 140) | (df['diastolic_bp'] >= 90), 'bp_category'] = 1  # 高血压
        df.loc[(df['systolic_bp'] < 90) | (df['diastolic_bp'] < 60), 'bp_category'] = -1  # 低血压
        
        # 血糖分类
        df['glucose_category'] = 0  # 正常
        df.loc[df['blood_glucose'] >= 126, 'glucose_category'] = 1  # 高血糖
        df.loc[df['blood_glucose'] < 70, 'glucose_category'] = -1  # 低血糖
        
        # 年龄分组
        df['age_group'] = pd.cut(df['age'], 
                                bins=[0, 30, 45, 60, 100], 
                                labels=[0, 1, 2, 3])  # 0: 青年, 1: 中年, 2: 中老年, 3: 老年
        
        # 风险评分计算
        df['risk_score'] = (
            df['age'] * 0.1 +
            df['bmi'] * 0.2 +
            df['systolic_bp'] * 0.15 +
            df['blood_glucose'] * 0.15 +
            df.get('smoking', 0) * 20 +
            df.get('family_history', 0) * 15 +
            df.get('exercise_frequency', 1) * -5
        )
        
        return df
    
    def _calculate_derived_features(self, data):
        """计算衍生特征"""
        # BMI分类
        bmi = data['bmi']
        if bmi < 18.5:
            data['bmi_category'] = 0  # 偏瘦
        elif bmi < 25:
            data['bmi_category'] = 1  # 正常
        elif bmi < 30:
            data['bmi_category'] = 2  # 超重
        else:
            data['bmi_category'] = 3  # 肥胖
        
        # 血压分类
        systolic = data['systolic_bp']
        diastolic = data['diastolic_bp']
        if systolic >= 140 or diastolic >= 90:
            data['bp_category'] = 1  # 高血压
        elif systolic < 90 or diastolic < 60:
            data['bp_category'] = -1  # 低血压
        else:
            data['bp_category'] = 0  # 正常
        
        # 血糖分类
        glucose = data['blood_glucose']
        if glucose >= 126:
            data['glucose_category'] = 1  # 高血糖
        elif glucose < 70:
            data['glucose_category'] = -1  # 低血糖
        else:
            data['glucose_category'] = 0  # 正常
        
        # 年龄分组
        age = data['age']
        if age < 30:
            data['age_group'] = 0  # 青年
        elif age < 45:
            data['age_group'] = 1  # 中年
        elif age < 60:
            data['age_group'] = 2  # 中老年
        else:
            data['age_group'] = 3  # 老年
        
        # 风险评分
        data['risk_score'] = (
            age * 0.1 +
            bmi * 0.2 +
            systolic * 0.15 +
            glucose * 0.15 +
            data.get('smoking', 0) * 20 +
            data.get('family_history', 0) * 15 +
            data.get('exercise_frequency', 1) * -5
        )
        
        return data
    
    def _fit_scaler(self, df):
        """训练标准化器"""
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        self.scaler.fit(df[numeric_columns])
        self.feature_columns = list(numeric_columns)
        self.is_fitted = True
    
    def generate_sample_data(self, n_samples=1000):
        """
        生成示例数据用于测试
        
        Args:
            n_samples (int): 样本数量
            
        Returns:
            DataFrame: 生成的示例数据
        """
        np.random.seed(42)
        
        data = {
            'age': np.random.normal(50, 15, n_samples).astype(int),
            'gender': np.random.choice([0, 1], n_samples),
            'height': np.random.normal(170, 10, n_samples).astype(int),
            'weight': np.random.normal(70, 15, n_samples).astype(int),
            'systolic_bp': np.random.normal(120, 20, n_samples).astype(int),
            'diastolic_bp': np.random.normal(80, 15, n_samples).astype(int),
            'blood_glucose': np.random.normal(100, 25, n_samples).astype(int),
            'cholesterol': np.random.normal(200, 40, n_samples).astype(int),
            'heart_rate': np.random.normal(72, 12, n_samples).astype(int),
            'smoking': np.random.choice([0, 1], n_samples, p=[0.7, 0.3]),
            'drinking': np.random.choice([0, 1], n_samples, p=[0.6, 0.4]),
            'family_history': np.random.choice([0, 1], n_samples, p=[0.8, 0.2]),
            'exercise_frequency': np.random.choice([0, 1, 2, 3], n_samples, p=[0.2, 0.3, 0.3, 0.2]),
            'diet_preference': np.random.choice([0, 1, 2], n_samples, p=[0.3, 0.5, 0.2])
        }
        
        df = pd.DataFrame(data)
        
        # 计算BMI
        df['bmi'] = df['weight'] / ((df['height'] / 100) ** 2)
        
        # 生成标签（模拟）
        df['hypertension_risk'] = np.random.choice([0, 1], n_samples, p=[0.7, 0.3])
        df['diabetes_risk'] = np.random.choice([0, 1], n_samples, p=[0.8, 0.2])
        
        return df
